#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import torch
import logging
import subprocess

# 设置日志级别为INFO
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 定义一个函数来运行nvidia-smi并获取特定进程的GPU使用情况
def get_process_gpu_info(pid):
    try:
        # 运行nvidia-smi命令并捕获输出
        result = subprocess.run([
            'nvidia-smi', 
            '--format=csv,noheader', 
            '--query-compute-apps=pid,gpu_name,gpu_uuid,used_memory'
        ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        
        # 解析输出
        for line in result.stdout.strip().split('\n'):
            if line.strip():
                parts = line.split(',')
                if len(parts) >= 4 and parts[0].strip() == str(pid):
                    return {
                        'pid': pid,
                        'gpu_name': parts[1].strip(),
                        'gpu_uuid': parts[2].strip(),
                        'used_memory': parts[3].strip()
                    }
        return None
    except Exception as e:
        logger.error(f"获取GPU进程信息失败: {str(e)}")
        return None

# 定义一个函数来运行gpustat并获取物理GPU的使用情况
def get_gpustat_info():
    try:
        # 运行gpustat命令并捕获输出
        result = subprocess.run(['gpustat'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        logger.info("\n=== 当前GPU使用情况 (gpustat) ===")
        for line in result.stdout.strip().split('\n'):
            logger.info(line)
    except Exception as e:
        logger.error(f"运行gpustat失败: {str(e)}")

logger.info("===== 验证CUDA设备映射机制 =====")

# 记录当前进程ID
current_pid = os.getpid()
logger.info(f"当前进程ID: {current_pid}")

# 1. 初始状态
logger.info("\n=== 初始状态 ===")
logger.info(f"初始 CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', '未设置')}")
logger.info(f"初始可用GPU数量: {torch.cuda.device_count() if torch.cuda.is_available() else 0}")

# 2. 设置CUDA_VISIBLE_DEVICES=1
logger.info("\n=== 设置CUDA_VISIBLE_DEVICES=1后 ===")
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
torch.cuda.empty_cache()
logger.info(f"设置后 CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES')}")
logger.info(f"设置后可用GPU数量: {torch.cuda.device_count()}")

# 3. 创建一个大张量来占用GPU内存，这样可以在nvidia-smi中看到
logger.info("\n=== 创建大张量并查看逻辑设备映射 ===")
# 创建一个占用约5GB内存的张量
tensor_size = (1024, 1024, 1024)  # 约4GB
large_tensor = torch.randn(*tensor_size, device='cuda')
logger.info(f"大张量创建在设备: {large_tensor.device}")

# 4. 查看当前进程在物理GPU上的实际位置
logger.info("\n=== 查看当前进程在物理GPU上的实际位置 ===")
logger.info("gpustat显示的是物理GPU编号，而我们的程序使用的是逻辑编号")
process_info = get_process_gpu_info(current_pid)
if process_info:
    logger.info(f"当前进程在物理GPU上的信息: {process_info}")
else:
    logger.warning("无法获取当前进程的GPU信息")

# 5. 显示gpustat输出以验证
get_gpustat_info()

# 6. 解释结果
logger.info("\n===== 结果解释 =====")
logger.info("1. 当设置CUDA_VISIBLE_DEVICES=1时，这个物理GPU在程序内部会被重新映射为'cuda:0'")
logger.info("2. 这是CUDA的标准行为 - CUDA会将可见的GPU重新编号为从0开始的连续编号")
logger.info("3. 因此，即使我们的程序报告模型加载到'cuda:0'，实际上它是加载到物理GPU 1上的")
logger.info("4. 这就是为什么gpustat显示物理GPU 0上有程序运行，但那是另一个进程")
logger.info("\n总结: 我们的GPU选择和模型加载逻辑工作正常，模型确实被加载到了我们选择的物理GPU 1上，只是在程序内部它被重新命名为'cuda:0'")

logger.info("\n验证完成！")